//////////////////////////////////////////////////////////////////////////////////
//	 Copyright 2011 by Yosef Grabivker, ryuho@homenetworksoft.com
//
//   Licensed under the Apache License, Version 2.0 (the "License");
//   you may not use this file except in compliance with the License.
//   You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
//   Unless required by applicable law or agreed to in writing, software
//   distributed under the License is distributed on an "AS IS" BASIS,
//   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//   See the License for the specific language governing permissions and
//   limitations under the License.
//////////////////////////////////////////////////////////////////////////////////

#include <wx/sstream.h>
#include <wx/protocol/http.h>
#include <wx/wx.h>
#include "telnetlink.h"
#include "sshsocket.h"

#define MAX_LOGIN_PROMPT_WAIT 5

CTelnetLink::CTelnetLink(void):m_bGrepSupported(false),
	m_chPrompt('#'), m_bLogedOn(false)
{
}

CTelnetLink::~CTelnetLink(void)
{
	CleanUp();
}

wxString CTelnetLink::GetEnterSeq()
{
	return m_sock->GetEnterSeq();
}

//not thread safe
void CTelnetLink::Create()
{
}

RTSTRING CTelnetLink::WaitForLoginString( const wxString& strLogin )
{
	wxString strTmp, strLoginLowCase;
	RTSTRING rtStr, rtStrRes;

	strLoginLowCase = strLogin.Lower();

	time_t tStart=time(NULL);

	do
	{
		rtStr = ReadString(1);
		rtStrRes+=rtStr;

		if ( !rtStr.empty() )
		{
			strTmp = wxString::FromAscii( rtStr.c_str() ).Lower();

			if ( wxNOT_FOUND!=strTmp.Find(strLogin) )
			{
				break;
			}
		}
	}while(time(NULL)-tStart<MAX_LOGIN_PROMPT_WAIT); 

	return rtStrRes;
}

//TODO: detect command prompt
void CTelnetLink::Connect(const CONPARAMS& con)
{
	const wxString 	strRouterAddr = con.strIP;
	const wxString  strUser = con.strUsr;
	const wxString	strPassw = con.strPassw;
	unsigned short wPort = con.wPort;
	cPROTOCOL proto = con.cProto;
	RTSTRING rtStr;

	if ( con.bHTTPEnable && !con.strHTTP.empty() )
	{
		EnableTelnetByHTTP(con);
	}

	m_sock.SetType(proto);
	m_sock->Connect(strRouterAddr.ToAscii(), wPort);

	if ( proto==cPROTOCOLTelnet )
	{
		//Wait for loging prompt
		if ( !con.strLoginPrompt.empty() )
		{
			rtStr = WaitForLoginString( con.strLoginPrompt );
		}
		else
		{
			rtStr = ReadString(1);

			if ( rtStr.empty() )
			{
				wxThread::Sleep(500);
				rtStr = ReadString(1);
			}
		}

		m_sock->Purge();

		wxString strWelcomePrompt( wxString::FromAscii(rtStr.c_str()).MakeLower() );
		bool bUserEmpty(true), bPassEmpty(true);
		bool bPasswordRequest(-1!=strWelcomePrompt.Find(wxT("password") ) );

		//Some routers, like Vigor asks password only
		//Send user name
		if ( !strUser.empty() && !bPasswordRequest )
		{
			SendString(strUser);
			SendEnter();
			wxThread::Sleep(250);
			ReadString(1);
			m_sock->Purge();
			bUserEmpty = false;
		}

		if ( !strPassw.empty() )
		{
			//Send password
			SendString(strPassw);
			SendEnter();
			
			bPassEmpty = false;
		}

		if ( bUserEmpty )
		{
			SendEnter();
		}

		if ( bPassEmpty )
		{
			SendEnter();
		}

		wxThread::Sleep(300);

		//Check common error messages
		if ( !WaitForPrompt() )
		{
			THROW_EXC(cERWrongUserOrPassword);
		}
	}
	else
	{
		m_sock->Login( con );
		m_sock->Purge();
	}

	if ( !con.strShell.empty() )
	{
		//Run shell
		SendString( con.strShell );
		SendEnter();
		wxThread::Sleep(200);
		WaitForPrompt();
	}

	m_sock->Purge();

	m_bLogedOn = true;
}


//TODO:redesign this
RTSTRING CTelnetLink::ReadFile(const wxString& strPath)
{
	RTSTRING strRes;

	if ( !m_sock->IsOk() )  THROW_EXC(cERWrongObjectState);
	if ( !m_sock->IsConnected() ) THROW_EXC(cERWrongObjectState);
	if ( strPath.empty() ) THROW_EXC(cERParam);
	
	m_sock->Purge();

	wxString strCMD(STR_CAT_CMD);
	strCMD+=_(" ");
	strCMD+=strPath;
	

	SendString( ( strCMD+GetEnterSeq() ).c_str() );

	int nCount(3);
	
	do
	{
		RTSTRING rtAnswer = ReadString();
		size_t nPos = rtAnswer.find( strCMD.ToAscii() );

		//echo received
		if ( std::string::npos!=nPos )
		{
			//cut echo
			rtAnswer.erase(0, nPos+strCMD.size() );
		}

		nPos = rtAnswer.rfind('\r');

		//cut prompt
		if ( std::string::npos!=nPos)
		{
			rtAnswer.erase(nPos, rtAnswer.size() );
		}

		size_t nPosNotExists = rtAnswer.find( "No such file or directory" );

		if ( std::string::npos!=nPosNotExists )
		{
			return "";
		}

		strRes=rtAnswer;

		#ifdef LINUX
			#ifdef DEBUG
			fprintf(stderr, "%s", strRes.c_str());
			#endif
		#else
		OutputDebugStringA(strRes.c_str());
		#endif
	}while( strRes.empty() && nCount-->0);

	/*
	if ( wxAnswer.StartsWith(strCMD+_(TELNET_ENTER), &strRes) )
	{
		strRes = strRes.BeforeLast('\r');
	}*/
	
	//printf("Res: %s\n", strRes.c_str());

	#ifdef DEBUG
	assert(!strRes.empty());
	#endif

	return strRes;
}

RTSTRING CTelnetLink::SendCmd(const wxString& strScript, unsigned int lDelay)
{
	RTSTRING strRes;

	if ( !m_sock->IsOk() )  THROW_EXC(cERWrongObjectState);
	if ( !m_sock->IsConnected() ) THROW_EXC(cERWrongObjectState);
	if ( strScript.empty() ) THROW_EXC(cERParam);

	wxString strCMD(strScript);
	strCMD+=GetEnterSeq();

	m_sock->Purge();

	SendString( strCMD.c_str() );

	RTSTRING rtAnswer = ReadString(-1, lDelay);

	size_t nPos = rtAnswer.find( strCMD.ToAscii() );

	if ( !nPos )
	{
		rtAnswer.erase(0, strCMD.size() );

		if ( rtAnswer.empty() )
			rtAnswer = ReadString(-1, lDelay);
	}
	else
	{
		size_t nPos = rtAnswer.find( TELNET_ENTER );

		try
		{
			rtAnswer.erase(0,  nPos==std::string::npos?strCMD.size():nPos );
		}
		catch(...)
		{
		}
	}

	nPos = rtAnswer.rfind('\r');

	if ( std::string::npos!=nPos)
	{
		rtAnswer.erase(nPos, rtAnswer.size() );
	}

	strRes=rtAnswer;

	return strRes;
}

RTSTRING CTelnetLink::SendSymb(const wxString& strScript, unsigned int lDelay)
{
	RTSTRING strRes;

	if ( !m_sock->IsOk() )  THROW_EXC(cERWrongObjectState);
	if ( !m_sock->IsConnected() ) THROW_EXC(cERWrongObjectState);
	if ( strScript.empty() ) THROW_EXC(cERParam);

	wxString strCMD(strScript);
	m_sock->Purge();

	SendString( strCMD.c_str() );

	strRes = ReadString(-1, lDelay);
	
	return strRes;
}


bool CTelnetLink::IsLoggedOn()
{
	return m_sock->IsConnected() && m_bLogedOn;
}

void CTelnetLink::CleanUp()
{
	if ( m_sock->IsOk() )
		m_sock->Close();

	m_sock.CleanUp();

	m_bLogedOn = false;
}

/////////////////////////////////////////////////////////////////////////////////////////
RTSTRING CTelnetLink::ReadString(wxUint32 lLinesExpected/*=-1*/, wxUint32 lTimeOut/*=RT_MAX_READ_TIMEOUT*/)
{
	char		buf[RT_BUF_SIZE];
	wxUint32	nChunkSize(sizeof(buf)-1);
	wxUint32	nReadTotal=0;
	int			nRead;
	
	std::string strTmp;

	if ( !m_sock->IsOk() ) THROW_EXC(cERWrongObjectState);
	if ( !m_sock->IsConnected() ) THROW_EXC(cERWrongObjectState);

	do
	{
		if ( !m_sock->IsConnected() ) break;

		memset(buf,0, sizeof(buf) );

		nRead = m_sock->Read( buf, nChunkSize);

		//probably connection closed by other side
		if ( !nRead )
		{
			break;
		}

		nRead = RemoveANSISequence(buf, nRead);

		buf[nRead] = '\0';

		#ifdef _DEBUG
			#ifdef LINUX
				fprintf(stderr, "%s", buf);
			#else
				OutputDebugStringA(buf);
			#endif
		#endif

		AddToLog(buf);

		nReadTotal+=nRead;
		strTmp+=buf;

		if ((wxUint32)-1!=lLinesExpected)
		{
			size_t		stFrom(-1);
			wxUint32	lCount(-1);

			do
			{
				lCount++;

				if ( lCount==lLinesExpected )
					return strTmp;

				stFrom=strTmp.find("\n", ++stFrom);

			}while(std::string::npos!=stFrom);
		}

		//#bug problem if prompt like simens ">"
		//command prompt found
		wxString str = wxString::FromAscii( strTmp.c_str() );
		str.Trim(true);

		size_t stPrompt = str.find(m_chPrompt);
		size_t stSize   = str.size();

		if ( wxNOT_FOUND!=stPrompt && stPrompt == stSize-1 )
		{
			break;
		}
	}
	while( m_sock->WaitForRead(lTimeOut) );

	return strTmp;
}

int	CTelnetLink::RemoveANSISequence(char* buf, int nRead)
{
	std::string str;
	
	char* ptrSeq;
	char* ptrSeqEnd = buf+nRead;

	bool bFound = false;

	for ( ptrSeq = buf;ptrSeq!=ptrSeqEnd;++ptrSeq )
	{
		if ( *ptrSeq==0x1b &&
			 *(ptrSeq+1)==0x5b ) //esc
		{
			ptrSeq++;
			bFound = true;
			continue;
		}

		if ( bFound && isalpha( *ptrSeq ) )
		{
			bFound = false;
			continue;
		}

		if ( bFound )
		{
			continue;
		}

		str+=*ptrSeq;
	}

	if ( !str.empty() )
	{
		std::copy(str.begin(), str.end(), buf );
	}

	return str.size();
}

void CTelnetLink::SendString(const wxString& strStr)
{
	if ( !m_sock->IsOk() ) THROW_EXC(cERWrongObjectState);
	if ( !m_sock->IsConnected() ) THROW_EXC(cERWrongObjectState);
	
	if ( !strStr.empty() )
		m_sock->Write( strStr.ToAscii(), strStr.size() );
	
	//fprintf(stderr, strStr.ToAscii());
}

void CTelnetLink::SendEnter()
{
	if ( !m_sock->IsOk() ) THROW_EXC(cERWrongObjectState);

	wxString strStr = GetEnterSeq();

	m_sock->Write( strStr.ToAscii(), strStr.size() );
}

bool CTelnetLink::IsPromptChar(char c)
{
	bool bRes(false);
	const char szPrompt[]="#>";

	for (size_t i=0;i<sizeof(szPrompt)/sizeof(szPrompt[0]);i++)
	{
		if ( szPrompt[i]==c )
		{
			bRes = true;
			break;
		}
	}

	return bRes;
}

bool CTelnetLink::WaitForPrompt()
{
	bool bRes(true);
	bool bPromptFound(false);

	RTSTRING rtStr;

	time_t tNow, tStart = time(NULL);

	do
	{
		rtStr = ReadString();
		rtStr = (const char*)wxString::FromAscii( rtStr.c_str() ).MakeLower().ToAscii();

		if ( std::string::npos!=rtStr.find(STR_WRONG_USER_OR_PASSW) )
		{
			bRes = false;
			break;
		}

		if ( std::string::npos!=rtStr.find("wrong") )
		{
			bRes = false;
			break;
		}
		
		if ( std::string::npos!=rtStr.find("incorrect") )
		{
				bRes = false;
				break;
		}

		if ( std::string::npos!=rtStr.find("invalid") )
		{
			bRes = false;
			break;
		}

		if ( std::string::npos!=rtStr.find("error") )
		{
			bRes = false;
			break;
		}

		if ( std::string::npos!=rtStr.find("failed") )
		{
			bRes = false;
			break;
		}

		if ( std::string::npos!=rtStr.find("login:") )
		{
			bRes = false;
			break;
		}

		if ( std::string::npos!=rtStr.find("password:") )
		{
			bRes = false;
			break;
		}

		size_t nPos = rtStr.rfind(TELNET_ENTER);

		//nPos may be 0 for Simiens router
		if ( /*nPos &&*/ std::string::npos!=nPos )
		{
			//Prompt may be like >ENTER in case of login script
			if ( nPos ) --nPos;

			for( size_t i=nPos;i<rtStr.size();i++ )
			{
				char cPrompt=rtStr[i];

				if ( IsPromptChar(cPrompt) )
				{
					if ( cPrompt!=']' &&
						 cPrompt!='[' &&
						 cPrompt!=')' &&
						 cPrompt!='(' )
					{
						bPromptFound = true;
						m_chPrompt = cPrompt;
						break;
					}
				}
			}
		}

		tNow = time(NULL);
	}while(!bPromptFound && tNow-tStart<MAX_LOGIN_TIMEOUT);

	return bRes;
}

void CTelnetLink::AddToLog(const char* szStr)
{
	try
	{
		if ( m_strLog.size()>0x1000 )
		{
			int nLines=100;

			while( nLines-->0 && m_strLog.size()>0x1000 )
			{
				size_t nPos = m_strLog.find('\n');

				if ( nPos!= std::string::npos )
				{
					m_strLog.erase( 0, nPos+1);
				}
			}

			if ( m_strLog.size()>0x1000 )
			{
				m_strLog.erase( m_strLog.begin(), m_strLog.end());
			}
		}

		m_strLog+=szStr;
	}
	catch(...)
	{
	}
}

/////////////////////////////////////////////////////////////////////////////////
CRTSocket* CTelnetLink::CSockGuard::CreateNewSocket()
{
	CRTSocket* pSocket(NULL);

	switch(m_proto)
	{
		case cPROTOCOLTelnet:
		{
			pSocket = new CRTelnetSocket();
		}
		break;

		#ifdef USE_SSH_SOCKET
		case cPROTOCOLSSH:
		{
			pSocket = new CRSSHSocket();
		}
		break;
		#endif
	}

	if ( !pSocket ) THROW_EXC(cERWrongObjectState);

	m_pSocket = pSocket;

	return m_pSocket;
}

void CTelnetLink::EnableTelnetByHTTP(const CONPARAMS& con)
{
	wxHTTP get;
	get.SetHeader(_T("Content-type"), _T("text/html; charset=utf-8"));
	get.SetTimeout(10);
	get.SetUser( con.strUsr );
	get.SetPassword( con.strPassw );

	// this will wait until the user connects to the internet. It is important in case of dialup (or ADSL) connections
	int nRet=3;
	while (nRet-->0 && !get.Connect(con.strIP))
		wxSleep(1);

	wxString strURL=con.strHTTP;

	wxInputStream *httpStream = NULL;
	
	do 
	{
		httpStream = get.GetInputStream(strURL);

		if (httpStream &&
			get.GetError() == wxPROTO_NOERR)
		{
			wxString res;
			wxStringOutputStream out_stream(&res);
			httpStream->Read(out_stream);
			break;
		}

	}while (!httpStream && nRet-->0);

	wxDELETE(httpStream);
	get.Close();
}